#!/usr/bin/env python

class Node ( object ):
	def __init__ ( self, key, data = None ):
		self.key = key
		self.data = data
		self.children = []
		self.parent = None

		self._is_root = False

	def is_root ( self ):
		return self._is_root
	
	def add ( self, node ):
		self.children.append ( node )

	def __repr__ ( self ):
		return "<Node: %s %s>" % ( self.key, self.data )

	def in_depth ( self ):
		for n in self.children:
			if n.children:
				for n2 in n.in_depth ():
					yield n2
			else:
				yield n

		yield self

	def traverse ( self ):
		yield self

		for n in self.children:
			if n.children:
				for n2 in n.traverse ():
					yield n2
			else:
				yield n

	def get_level ( self, level ):
		assert level >= 0
		for n in self.children:
			if level == 1:
				yield n
			else:
				for n2 in n.get_level ( level - 1 ):
					yield n2

	def to_object ( self ):
		dct = dict ( key = self.key, data = self.data, children = [] )
		if self.children:
			for n in self.children:
				dct [ 'children' ].append ( n.to_object () )
		return dct
				

class Tree ( object ):
	def __init__ ( self ):
		self._nodes = {}
		self.root = Node ( '__ROOT__' )
		self.root._is_root = True

	def add ( self, key, data = None, parent = None ):
		n = Node ( key, data )
		self.node_add ( n, parent )

		return n

	def node_add ( self, node, parent = None ):
		if not parent:
			parent = self.root

		parent.add ( node )
		node.parent = parent

		self._nodes [ node.key ] = node

	def __getitem__ ( self, key ):
		return self._nodes.get ( key, None )

	def in_depth ( self, node = None ):
		if not node: node = self.root

		return node.in_depth ()

	def traverse ( self, node = None ):
		if not node: node = self.root

		return node.traverse ()


if __name__ == '__main__':
	tree = Tree ()
	tree.add ( 1, 'Fabio' )
	tree.add ( 2, 'Keeper', parent = tree [ 1 ] )
	tree.add ( 3, 'Bundle', parent = tree [ 1 ] )
	tree.add ( 4, 'Pig', parent = tree [ 2 ] )
	tree.add ( 5, 'Paola', parent = tree [ 3 ] )
	tree.add ( 6, 'Alby', parent = tree [ 3 ] )

	print [ x.data for x in tree.traverse ( tree [ 2 ] ) ]
